#include "grate/dfa.h"

#include <sstream>
#include <string>

namespace growth_rate {

bool DFAInterface::equals(const DFAInterface& other) const {
  if (size_ != other.size()) return false;
  if (alphabet_size_ != other.alphabet_size()) return false;
  for (size_t node = StartNode; node < size_; node++) {
    for (uint8_t letter = 0; letter < alphabet_size_; letter++) {
      if (get_next(node, letter) != other.get_next(node, letter))
        return false;
    }
  }
  return true;
}

std::string DFAInterface::dump_transitions() const {
  std::stringstream out;
  out << "[";
  for (size_t node = StartNode; node < size_; node++) {
    if (static_cast<int>(node) > StartNode) out << ",\n";
    out << "(";
    for (uint8_t letter = 0; letter < alphabet_size_; letter++) {
      if (letter > 0) out << ", ";
      out << get_next(node, letter);
    }
    out << ")";
  }
  out << "]";
  return out.str();
}

void DFAInterface::dump_valid_edges(DFAInterface::Edge *edges) const {
  size_t pos = 0;
  for (int node = StartNode; node < size_; node++) {
    for (uint8_t letter = 0; letter < alphabet_size_; letter++) {
      int next_node = get_next(node, letter);
      if (next_node >= StartNode) {
        edges[pos].from = node;
        edges[pos].to = next_node;
        pos++;
      }
    }
  }
}

LexminTrie::LexminTrie(size_t capacity, uint8_t alphabet_size)
    : DFAInterface(0, alphabet_size),
      capacity_(capacity),
      leaf_nodes_(0) {
  CHECK_GT(capacity_, 0);
  CHECK_GT(alphabet_size_, 1);
  edges_ = AllocateArray<int>(capacity_ * alphabet_size_);
  // Create fake zero node and start node.
  CHECK(create_empty_vertex());
  CHECK_EQ(1, size_);
}

bool LexminTrie::add_lexmin_word(const Word &word) {
  int cur_node = StartNode;
  size_t length = word.length();
  CHECK_GT(length, 0);
  uint8_t max_allowed_letter = 0;
  for (size_t i = 0; i < length - 1; i++) {
    uint8_t letter = word[i];
    CHECK_LT(letter, alphabet_size_);
    CHECK_LE(letter, max_allowed_letter);
    max_allowed_letter = Max(max_allowed_letter, (uint8_t)(letter + 1));
    int next_node = get_next(cur_node, letter);
    if (next_node == UnknownNode) {
      if (!create_empty_vertex()) return false;
      next_node = size_ - 1;
      set_next(cur_node, letter, next_node);
    } else if (next_node == FakeSink) {
      return false;
    }
    CHECK_GE(next_node, StartNode);
    cur_node = next_node;
  }
  int final_node = get_next(cur_node, word[length - 1]);
  CHECK(final_node == UnknownNode || final_node == FakeSink);
  set_next(cur_node, word[length - 1], FakeSink);
  leaf_nodes_++;
  return true;
}

bool LexminTrie::accepts_word(const Word &word, Direction direction) const {
  size_t length = word.length();
  if (length == 0) return true;
  int cur_node = StartNode;
  size_t first_idx, last_idx;
  if (direction == Forward) {
    first_idx = 0;
    last_idx = length - 1;
  } else {
    first_idx = length - 1;
    last_idx = 0;
  }
  // Build a map to set the way we should rename letters to get a lexmin
  // equivalent.
  uint8_t letters_map[256];
  uint8_t next_letter_to_map = 0;
  memset(letters_map, 0xFF, alphabet_size_);
  for (;; first_idx += direction) {
    uint8_t letter = word[first_idx];
    CHECK_LT(letter, alphabet_size_);
    if (letters_map[letter] == 0xFF) {
      letters_map[letter] = next_letter_to_map++;
    }
    letter = letters_map[letter];
    int next_node = get_next(cur_node, letter);
    if (next_node == UnknownNode) return true;
    if (next_node == FakeSink) return false;
    cur_node = next_node;
    if (first_idx == last_idx) break;
  }
  return true;
}

bool LexminTrie::create_empty_vertex() {
  if (size_ == capacity_)
    return false;
  CHECK_LT(size_, capacity_);
  for (uint8_t letter = 0; letter < alphabet_size_; letter++) {
    edges_[size_ * alphabet_size_ + letter] = UnknownNode;
  }
  size_++;
  return true;
}

}  // namespace growth_rate
